package langchain

import (
	"fmt"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
	"golang.org/x/net/context"
)

type OpenAI struct {
	modelPath string
	token     string
}

func NewOpenAI(url, token string) (*OpenAI, error) {

	if token == "" {
		return nil, fmt.Errorf("OpenAI token is required")
	}
	return &OpenAI{
		modelPath: url,
		token:     token,
	}, nil
}

func (s *OpenAI) PredictOpenAI(text string, opts ...PredictOption) (*Predict, error) {
	po := NewPredictOptions(opts...)
	llm, err := openai.New(openai.WithToken(s.token))
	if err != nil {
		return nil, err
	}
	co := []llms.CallOption{
		llms.WithModel(po.Model),
		llms.WithMaxTokens(po.MaxTokens),
		llms.WithTemperature(po.Temperature),
		llms.WithStopWords(po.StopWords)}
	ctx := context.Background()
	completion, err := llm.Call(ctx, text, co...)
	if err != nil {
		return nil, err
	}
	return &Predict{
		Completion: completion,
	}, nil
}
